import os
import cv2
import xml.etree.ElementTree as ET
from sklearn.model_selection import train_test_split
import shutil
from tqdm import tqdm
import glob

def convert_detrac_to_yolo(xml_dir, img_dir, output_dir, classes):
    # 清空并重新创建输出目录
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(os.path.join(output_dir, "images", "train"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "images", "val"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "labels", "train"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "labels", "val"), exist_ok=True)

    # 获取所有XML文件
    xml_files = sorted([f for f in os.listdir(xml_dir) if f.endswith('.xml')])
    print(f"共找到 {len(xml_files)} 个XML文件")

    # 分割训练集和验证集
    train_files, val_files = train_test_split(xml_files, test_size=0.2, random_state=42)
    print(f"训练集: {len(train_files)} 个文件, 验证集: {len(val_files)} 个文件")

    # 处理训练集
    print("\n处理训练集...")
    for xml_file in tqdm(train_files):
        process_xml_file(os.path.join(xml_dir, xml_file), img_dir, output_dir, "train", classes)

    # 处理验证集
    print("\n处理验证集...")
    for xml_file in tqdm(val_files):
        process_xml_file(os.path.join(xml_dir, xml_file), img_dir, output_dir, "val", classes)

    # 统计结果
    print("\n转换结果统计:")
    for split in ["train", "val"]:
        img_count = len(os.listdir(os.path.join(output_dir, "images", split)))
        label_count = len(os.listdir(os.path.join(output_dir, "labels", split)))
        print(f"{split}图片: {img_count}, {split}标签: {label_count}")

def process_xml_file(xml_path, img_dir, output_dir, split, classes):
    try:
        tree = ET.parse(xml_path)
        root = tree.getroot()
        seq_name = os.path.splitext(os.path.basename(xml_path))[0]
        width, height = 960, 540

        for frame in root.findall('.//frame'):
            frame_num = frame.get('num', '0')
            frame_num = str(int(frame_num)).zfill(5)  # 补齐五位

            img_name = f"img{frame_num}.jpg"

            # 递归查找所有包含seq_name的子目录下的图片
            possible_img_paths = glob.glob(os.path.join(img_dir, '**', seq_name, img_name), recursive=True)

            if not possible_img_paths:
                print(f"未找到图片: {img_name} in seq: {seq_name}")
                continue

            img_path = possible_img_paths[0]  # 只取第一个

            label_name = f"{seq_name}_{frame_num}.txt"
            label_path = os.path.join(output_dir, "labels", split, label_name)
            with open(label_path, 'w') as f:
                for target in frame.findall('.//target'):
                    box = target.find('box')
                    if box is None:
                        continue
                    xmin = float(box.get('left'))
                    ymin = float(box.get('top'))
                    w = float(box.get('width'))
                    h = float(box.get('height'))
                    x_center = (xmin + w / 2) / width
                    y_center = (ymin + h / 2) / height
                    w_norm = w / width
                    h_norm = h / height
                    f.write(f"{classes['vehicle']} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}\n")

            output_img_path = os.path.join(output_dir, "images", split, f"{seq_name}_{frame_num}.jpg")
            if not os.path.exists(output_img_path):
                try:
                    img = cv2.imread(img_path)
                    if img is not None:
                        cv2.imwrite(output_img_path, img)
                except Exception as e:
                    print(f"无法保存图片 {output_img_path}: {str(e)}")
    except Exception as e:
        print(f"处理文件 {xml_path} 时出错: {str(e)}")

if __name__ == "__main__":
    # 配置参数
    classes = {"vehicle": 0}
    xml_dir = "UA-DETRAC/DETRAC-Train-Annotations-XML"
    img_dir = "UA-DETRAC/DETRAC-train-data"
    output_dir = "yolo_data"

    # 检查路径是否存在
    if not os.path.exists(xml_dir):
        raise FileNotFoundError(f"XML目录不存在: {xml_dir}")
    if not os.path.exists(img_dir):
        raise FileNotFoundError(f"图片目录不存在: {img_dir}")

    print("开始转换数据集...")
    convert_detrac_to_yolo(xml_dir, img_dir, output_dir, classes)